import torch
import torch.nn as nn
import torch.nn.functional as F


class CIFAR10CNN(nn.Module):
    def __init__(self, num_classes=10):
        super(CIFAR10CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)

        # 计算全连接层输入尺寸
        self.fc_input_size = self._get_fc_input_size()
        self.fc1 = nn.Linear(self.fc_input_size, 512)
        self.fc2 = nn.Linear(512, num_classes)

    def _get_fc_input_size(self):
        """自动计算全连接层输入尺寸"""
        with torch.no_grad():
            x = torch.zeros(1, 3, 32, 32)  # CIFAR-10输入尺寸
            x = self.pool(F.relu(self.conv2(F.relu(self.conv1(x)))))
            return x.view(-1).shape[0]

    def forward(self, x):
        # 卷积层
        x = F.relu(self.conv1(x))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.dropout1(x)

        # 展平
        x = x.view(x.size(0), -1)  # 保持batch_size维度

        # 全连接层
        x = F.relu(self.fc1(x))
        x = self.dropout2(x)
        x = self.fc2(x)
        return x


# 测试模型
# model = CIFAR10CNN()
# test_input = torch.randn(128, 3, 32, 32)  # batch_size=128
# output = model(test_input)
# print(output.shape)  # 应该输出 torch.Size([128, 10])
